package cc.yaomm.gateway.supply.filter;

import cc.yaomm.common.basic.constant.DefaultConstant;
import io.netty.buffer.EmptyByteBuf;
import lombok.extern.slf4j.Slf4j;
import org.springframework.cloud.gateway.filter.GatewayFilterChain;
import org.springframework.cloud.gateway.filter.GlobalFilter;
import org.springframework.core.Ordered;
import org.springframework.core.io.buffer.DataBuffer;
import org.springframework.core.io.buffer.DataBufferUtils;
import org.springframework.core.io.buffer.NettyDataBufferFactory;
import org.springframework.http.HttpHeaders;
import org.springframework.http.MediaType;
import org.springframework.http.server.reactive.ServerHttpRequest;
import org.springframework.http.server.reactive.ServerHttpRequestDecorator;
import org.springframework.http.server.reactive.ServerHttpResponse;
import org.springframework.stereotype.Component;
import org.springframework.util.StringUtils;
import org.springframework.web.server.ServerWebExchange;
import reactor.core.publisher.Flux;
import reactor.core.publisher.Mono;

import java.nio.charset.StandardCharsets;

/**
 * <p>@Description : 此拦截器主要提取协议为：application/json的body数据 </p>
 * <p>@Author : QiLin.Xing </p>
 * <p>@Date : 2021/5/27 15:32 </p>
 */
@Slf4j
@Component
public class BodyFilter implements GlobalFilter, Ordered {

    @Override
    public Mono<Void> filter(ServerWebExchange exchange, GatewayFilterChain chain) {
        ServerHttpRequest request = exchange.getRequest();
        //获取协议
        String contentType = request.getHeaders().getFirst(HttpHeaders.CONTENT_TYPE);
        //判断协议不为空并且为application/json
        if (StringUtils.hasText(contentType) && contentType.contains(MediaType.APPLICATION_JSON_VALUE)) {
            ServerHttpResponse response = exchange.getResponse();
            //获取dataBuffer的创建工厂
            NettyDataBufferFactory factory = (NettyDataBufferFactory) response.bufferFactory();
            //解析body，并转发请求
            return DataBufferUtils.join(request.getBody())
                    //处理body为空的情况
                    .defaultIfEmpty(factory.wrap(new EmptyByteBuf(factory.getByteBufAllocator())))
                    .flatMap(dataBuffer -> {
                        byte[] bytes = new byte[dataBuffer.readableByteCount()];
                        dataBuffer.read(bytes);
                        String bodyString = new String(bytes, StandardCharsets.UTF_8);
                        //将解析到body内容缓存到请求上下文中
                        exchange.getAttributes().put(DefaultConstant.REQUEST_BODY, bodyString);
                        //必须释放掉，否则存在内存溢出的问题
                        DataBufferUtils.release(dataBuffer);
                        Flux<DataBuffer> cachedFlux = Flux.defer(() -> {
                            DataBuffer buffer = exchange.getResponse().bufferFactory().wrap(bytes);
                            return Mono.just(buffer);
                        });

                        //转发请求
                        ServerHttpRequest mutatedRequest = new ServerHttpRequestDecorator(
                                exchange.getRequest()) {
                            @Override
                            public Flux<DataBuffer> getBody() {
                                return cachedFlux;
                            }
                        };
                        return chain.filter(exchange.mutate().request(mutatedRequest).build());
                    });
        } else {
            return chain.filter(exchange);
        }
    }

    @Override
    public int getOrder() {
        return -200;
    }
}
